import torch
import torch.nn as nn


class unetConv2(nn.Module):
    def __init__(self, in_size, out_size, is_batchnorm):
        super(unetConv2, self).__init__()

        if is_batchnorm:
            self.conv1 = nn.Sequential(
                nn.Conv2d(in_size, out_size, 3, 1, padding=1, bias=True),
                nn.BatchNorm2d(out_size),
                nn.LeakyReLU(),
            )
            self.conv2 = nn.Sequential(
                nn.Conv2d(out_size, out_size, 3, 1, padding=1, bias=True),
                nn.BatchNorm2d(out_size),
                nn.LeakyReLU(),
            )
        else:
            self.conv1 = nn.Sequential(
                nn.Conv2d(in_size, out_size, 3, 1, padding=1, bias=True), nn.LeakyReLU()
            )
            self.conv2 = nn.Sequential(
                nn.Conv2d(out_size, out_size, 3, 1, padding=1, bias=True),
                nn.LeakyReLU(),
            )

    def forward(self, inputs):
        outputs = self.conv1(inputs)
        outputs = self.conv2(outputs)
        return outputs


class unetUp(nn.Module):
    def __init__(self, in_size, out_size, is_deconv):
        super(unetUp, self).__init__()
        if is_deconv:
            self.up = nn.Sequential(
                nn.ConvTranspose2d(
                    in_size,
                    out_size,
                    kernel_size=3,
                    stride=2,
                    padding=1,
                    output_padding=1,
                ),
                nn.LeakyReLU(),
            )
        else:
            self.up = nn.UpsamplingBilinear2d(scale_factor=2)

    def forward(self, inputs1, inputs2):
        outputs2 = self.up(inputs2)
        return torch.cat([inputs1, outputs2], 1)


class unet_v1(nn.Module):
    def __init__(
        self,
        in_channels: int = 1,
        out_channels: int = 1,
        is_deconv: bool = True,
        is_batchnorm: bool = False,
    ):
        super(unet_v1, self).__init__()
        self.is_deconv = is_deconv
        self.in_channels = in_channels
        self.is_batchnorm = is_batchnorm

        filters = [32, 64, 128, 256, 512]

        # downsampling
        self.conv1 = unetConv2(self.in_channels, filters[0], self.is_batchnorm)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True)

        self.conv2 = unetConv2(filters[0], filters[1], self.is_batchnorm)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True)

        self.conv3 = unetConv2(filters[1], filters[2], self.is_batchnorm)
        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True)

        self.conv4 = unetConv2(filters[2], filters[3], self.is_batchnorm)
        self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True)

        self.center = unetConv2(filters[3], filters[4], self.is_batchnorm)

        # upsampling
        self.up_concat4 = unetUp(filters[4], filters[3], self.is_deconv)
        self.conv5 = unetConv2(filters[3] * 2, filters[3], self.is_batchnorm)

        self.up_concat3 = unetUp(filters[3], filters[2], self.is_deconv)
        self.conv6 = unetConv2(filters[2] * 2, filters[2], self.is_batchnorm)

        self.up_concat2 = unetUp(filters[2], filters[1], self.is_deconv)
        self.conv7 = unetConv2(filters[1] * 2, filters[1], self.is_batchnorm)

        self.up_concat1 = unetUp(filters[1], filters[0], self.is_deconv)
        self.conv8 = unetConv2(filters[0] * 2, filters[0], self.is_batchnorm)

        # final conv
        self.final = nn.Conv2d(filters[0], out_channels, 3, 1, 1, bias=True)

    def forward(self, inputs):
        conv1 = self.conv1(inputs)
        maxpool1 = self.maxpool1(conv1)
        conv2 = self.conv2(maxpool1)
        maxpool2 = self.maxpool2(conv2)
        conv3 = self.conv3(maxpool2)
        maxpool3 = self.maxpool3(conv3)
        conv4 = self.conv4(maxpool3)
        maxpool4 = self.maxpool4(conv4)

        center = self.center(maxpool4)

        up4 = self.up_concat4(conv4, center)
        conv5 = self.conv5(up4)
        up3 = self.up_concat3(conv3, conv5)
        conv6 = self.conv6(up3)
        up2 = self.up_concat2(conv2, conv6)
        conv7 = self.conv7(up2)
        up1 = self.up_concat1(conv1, conv7)
        conv8 = self.conv8(up1)

        final = self.final(conv8)

        out = torch.add(final, inputs)

        return out

from torch.profiler import profile, record_function, ProfilerActivity

model = unet_v1()
inputs = torch.randn(1, 1, 640, 480)
from ptflops import get_model_complexity_info
macs, params = get_model_complexity_info(
    model, (1, 640, 480),
    as_strings=True,  # 以字符串形式输出
    print_per_layer_stat=True,  # 打印每层的统计信息
    verbose=True
)
print(f"计算量: {macs}")
print(f"参数量: {params}")
